
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;


/**
 * 热更新类加载器，支持在运行时动态替换已加载的类
 * 通过缓存机制和备份机制实现类的热替换，避免应用重启
 */
public class HotSwapClassLoader extends ClassLoader {

    // 已加载类的缓存（线程安全Map，避免重复加载）
    private Map<String, Class<?>> classCache = new ConcurrentHashMap<>();
    // 类字节码缓存（存储已加载类的原始字节码，用于回滚）
    private Map<String, byte[]> classBytes = new ConcurrentHashMap<>();
    // 类文件存储目录（如：/path/to/classes,d:/classes等）
    private String classDir;
    // 类版本备份（存储每个类的当前版本，用于回滚）
    private Map<String, ClassLoader> backupLoaders = new ConcurrentHashMap<>();

    public HotSwapClassLoader(String classDir) {
        this.classDir = classDir;
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        // 1. 优先从缓存中获取已加载的类（避免重复加载）
        if (classCache.containsKey(name)) {
            return classCache.get(name);
        }

        try {
            // 2. 从文件系统加载类字节码
            byte[] classData = loadClassData(name);
            if (classData == null) {
                throw new ClassNotFoundException("Class not found: " + name);
            }
            // 3. 定义新类（JVM级加载）
            Class<?> clazz = defineClass(name, classData, 0, classData.length);
            // 4. 缓存类和字节码（供后续热更新和回滚使用）
            classCache.put(name, clazz);
            classBytes.put(name, classData);
            return clazz;
        } catch (IOException e) {
            throw new ClassNotFoundException("Failed to load class: " + name, e);
        }
    }

    /**
     * 从文件系统加载类的字节码
     * @param className 类全限定名（如：com.example.MyClass）
     * @return 类字节码，不存在返回null
     */
    private byte[] loadClassData(String className) throws IOException {
        // 构建类文件路径（将包路径转为文件系统路径）
        String path = className.replace('.', File.separatorChar) + ".class";
        File file = new File(classDir, path);

        if (!file.exists()) {
            return null;
        }
        // 安全读取文件（try-with-resources确保流关闭）
        try (FileInputStream fis = new FileInputStream(file);
             ByteArrayOutputStream baos = new ByteArrayOutputStream()) {

            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = fis.read(buffer)) != -1) {
                baos.write(buffer, 0, bytesRead);
            }
            return baos.toByteArray();
        }
    }

    /**
     * 执行热更新：替换已加载的类
     * @param className 要热更新的类全限定名
     * @return true成功，false失败
     */
    public synchronized boolean hotSwapClass(String className) {
        try {
            // 1. 备份当前类版本（关键：避免更新失败时无法回滚）
            backupCurrentVersion(className);

            // 2. 清除缓存（触发重新加载）
            classCache.remove(className);
            classBytes.remove(className);

            // 3. 重新加载类（通过findClass触发）
            Class<?> newClass = findClass(className);
            if (newClass != null) {
                System.out.println("Successfully hot-swapped class: " + className);
                return true;
            }
        } catch (Exception e) {
            System.err.println("Failed to hot-swap class: " + className);
            e.printStackTrace();
            // 4. 失败时回滚到备份版本
            rollbackClass(className);
        }
        return false;
    }


    /**
     * 备份当前类的版本（为回滚做准备）
     * @param className 要备份的类
     */
    private void backupCurrentVersion(String className) {
        // 仅当类已加载时才备份
        if (classCache.containsKey(className)) {
            // 创建独立的类加载器保存备份（避免与主加载器冲突）
            BackupClassLoader backupLoader = new BackupClassLoader(
                    classBytes.get(className), // 备份字节码
                    classCache.get(className), // 备份类对象
                    this // 父加载器（保持类加载链）
            );
            backupLoaders.put(className, backupLoader);
        }
    }

    /**
     * 回滚到类的备份版本
     * @param className 要回滚的类
     * @return true成功，false失败
     */
    public synchronized boolean rollbackClass(String className) {
        ClassLoader backupLoader = backupLoaders.get(className);
        if (backupLoader instanceof BackupClassLoader) {
            BackupClassLoader backup = (BackupClassLoader) backupLoader;
            try {
                // 1. 从备份中恢复类和字节码
                classCache.put(className, backup.getBackupClass());
                classBytes.put(className, backup.getBackupBytes());
                System.out.println("Rolled back class: " + className);
                return true;
            } catch (Exception e) {
                System.err.println("Failed to rollback class: " + className);
                e.printStackTrace();
            }
        }
        return false;
    }

    /**
     * 获取所有已加载的类名
     * @return 类名集合（线程安全副本）
     */
    public Set<String> getLoadedClasses() {
        return new HashSet<>(classCache.keySet());
    }

    /**
     * 清理备份（防止内存泄漏，通常在不再需要热更新时调用）
     */
    public void cleanupBackups() {
        backupLoaders.clear();
    }


    /**
     * 备份类加载器：独立保存类的字节码和类对象
     * 通过继承ClassLoader实现，但不用于加载新类（仅存储备份）
     */
    private static class BackupClassLoader extends ClassLoader {
        private final byte[] backupBytes; // 备份的类字节码
        private final Class<?> backupClass; // 备份的Class对象


        public BackupClassLoader(byte[] bytes, Class<?> clazz, ClassLoader parent) {
            super(parent); // 保持类加载链
            this.backupBytes = bytes;
            this.backupClass = clazz;
        }

        // 仅提供备份数据的访问方法
        public byte[] getBackupBytes() {
            return backupBytes;
        }

        public Class<?> getBackupClass() {
            return backupClass;
        }
    }
}